# Copyright (C) 2022 ByteDance Inc.
# All rights reserved.
# Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).

# The software is made available under Creative Commons BY-NC-SA 4.0 license
# by ByteDance Inc. You can use, redistribute, and adapt it
# for non-commercial purposes, as long as you (a) give appropriate credit
# by citing our paper, (b) indicate any changes that you've made,
# and (c) distribute any derivative works under the same license.

# THE AUTHORS DISCLAIM ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING ALL
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR ANY PARTICULAR PURPOSE.
# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL
# DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
# WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING
# OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

import os
import argparse
import shutil
from einops import einops
import numpy as np
import imageio
import time
import torch
import sys, os

from tqdm import tqdm

sys.path.insert(0, os.getcwd())
from models import make_model
from visualize.utils import (
    generate,
    safe_mkdir,
    tensor2image,
    tensor2seg,
    get_truncation_mean,
)


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--ckpt", type=str, required=True, help="model checkpoint path")
    parser.add_argument(
        "--outdir",
        type=str,
        default="./results/samples/",
        help="path to the output directory",
    )
    parser.add_argument(
        "--sample_cnt", type=int, default=16, help="sample count for generation"
    )
    parser.add_argument(
        "--sample_batch", type=int, default=16, help="batch size for generation"
    )
    parser.add_argument(
        "--truncation", type=float, default=0.7, help="truncation ratio"
    )
    parser.add_argument(
        "--truncation_cnt",
        type=int,
        default=10000,
        help="number of vectors to calculate mean for the truncation",
    )
    parser.add_argument(
        "--truncation_batch",
        type=int,
        default=20,
        help="batch to calculate truncation mean",
    )

    parser.add_argument(
        "--save_latent",
        action="store_true",
        help="whether to save the output latent codes",
    )
    parser.add_argument(
        "--device", type=str, default="cuda:0", help="running device for inference"
    )
    args = parser.parse_args()

    safe_mkdir(args.outdir)

    print("Loading model ...")
    ckpt = torch.load(args.ckpt, map_location=args.device)
    ckpt["args"].batch = args.sample_batch
    model = make_model(ckpt["args"])
    model.to(args.device)
    model.eval()
    model.load_state_dict(ckpt["g_ema"])

    print("Generating images ...")
    start_time = time.time()
    # calculate the truncation means for each parts
    truncation_means = get_truncation_mean(
        model, args.truncation_cnt, args.truncation_batch, args.device
    )
    with torch.no_grad():
        loop_cnt = args.sample_cnt // args.sample_batch
        pbar = tqdm(range(loop_cnt), dynamic_ncols=True, smoothing=0.01)
        for step in pbar:
            images, segs, styles, _, querys = model(
                truncation=args.truncation,
                return_latents=True,
                truncation_latent=truncation_means,
            )
            images = tensor2image(images)
            segs = tensor2seg(segs)
            # change list querys to array
            querys = [q.cpu().numpy() for q in querys]
            querys = np.stack(querys, axis=0)  # n 1 b c
            for i in range(len(images)):
                imageio.imwrite(
                    f"{args.outdir}/{str(step*len(images)+i).zfill(6)}_img.jpg",
                    images[i],
                )
                imageio.imwrite(
                    f"{args.outdir}/{str(step*len(images)+i).zfill(6)}_seg.jpg", segs[i]
                )
                if args.save_latent:
                    np.save(
                        f"{args.outdir}/{str(step*len(images)+i).zfill(6)}_latent.npy",
                        styles[i : i + 1].cpu().numpy(),
                    )
                    np.save(
                        f"{args.outdir}/{str(step*len(images)+i).zfill(6)}_query.npy",
                        querys[:, :, i : i + 1],
                    )
    print(f"Average speed: {(time.time() - start_time)/(args.sample_cnt)}s")
